from pymetasploit3.msfrpc import MsfRpcClient
from mcp.server.fastmcp import FastMCP
# Initialize FastMCP server
mcp = FastMCP("Mytest")

client = None
shell = None
@mcp.tool()
def init_clinet() -> any:
    """
    在调用generate_console_bycid函数之前所必须运行的函数

    Args:
    
    return:
        返回clinet创建是否成功
    """
    global client
    client = MsfRpcClient('123',port=55552)
    return True

# Destroy a console
#client.console.console('1').destroy
@mcp.tool()
def run_command(cmd: str,cid:str) -> str:
    """
    这是渗透测试工具metasploitframework的交互式命令行
    输入的命令确保是一行命令完成之后就不再使用该会话的

    Args:
        cmd: 运行metasploitframework命令的内容，确保内容是一次性就能达到效果的。
        cid: 在使用该函数之前使用generate_console_bycid函数生成的会话cid
        
    return:
        返回命令执行结果
    """
    # Write to console
    client.consoles.console(cid).write(cmd)
    while True:
        if not client.consoles.console(cid).is_busy():
            break

    readcontent = client.consoles.console(cid).read()
    client.consoles.console(cid).destroy
    return readcontent

@mcp.tool()
def run_command_oneline(cmd: str,cid:str) -> str:
    """
    这是渗透测试工具metasploitframework的交互式命令行
    在这个函数中确保是多次使用相同会话的场景下使用的，相同会话的判断通过函数返回的cid作为依据
    比如使用某个模块之后再次调用该函数来设置参数这种情况

    Args:
        cmd: 运行metasploitframework命令的内容，确保内容是使用在多次调用该函数才能达到效果的
        cid: 在使用该函数之前使用generate_console_bycid函数生成的会话cid
        
    return:
        readcontent:返回命令执行结果
        cid: 当前console的会话cid
    """
    # Write to console
    client.consoles.console(cid).write(cmd)
    while True:
        if not client.consoles.console(cid).is_busy():
            break

    readcontent = client.consoles.console(cid).read()
    return readcontent,cid

@mcp.tool()
def destroy_console_bycid(cid: str) -> str:
    """
    在使用run_command_oneline函数之后在确保用户需求完全达成之后删除该console会话

    Args:
        
        cid: 需要删除的会话id
    return:
        返回命令执行结果
    """
    client.consoles.console(cid).destroy
    return "函数已执行"

@mcp.tool()
def generate_console_bycid() -> str:
    """
    在使用run_command_oneline函数或者run_command函数之前在确保用户输入任何msf的命令之前生成该console会话
    只要调用Mytest这个mcp，这是最优先调用的函数
    确保在调用该函数之前已经调用过init_clinet函数并获取到client对象
    Args:
        
    return:
        cid:返回生成的会话cid
    """
    
    cid = client.consoles.console().cid
    return cid

@mcp.tool()
def list_all_shell()->list:
    """
    在该函数调用之前请确保已经运行过init_clinet这个函数
    当用户询问到msf中的session管理相关或者后渗透相关的问题的时候优先调用该函数
    这个函数会返回当前所有可用的shell的id

    Args:
        
    return:
        shell_list:当前所有可用的shell的id
    """
    # Find all available sessions
    #print("Sessions avaiables : ")
    shell_list=[]
    for s in client.sessions.list.keys():
        shell_list.append(s)
    return shell_list

@mcp.tool()
def get_a_shell(index:int)->bool:
    """
    在该函数调用之前请确保已经运行过list_all_shell这个函数
    该函数会返回一个具体的shell对象供用户做后续处理


    Args:
        index:选择第几个shell(从0开始计数)
    return:
        True:表示函数运行过
    """
    global shell
    shell = client.sessions.session(list(client.sessions.list.keys())[index])
    return True

@mcp.tool()
def write_to_shell(cmd:str)->str:
    """
    在该函数调用之前请确保已经运行过get_a_shell这个函数
    该函数的输入内容是已经处在session中的命令，不支持“bg”这个命令
    该函数会返回命令执行的结果
    该函数在stop_shell调用前可以重复调用write_to_shell多次

    Args:
        cmd:已经处在session中的命令，不支持“bg”这个命令
    return:
        readcontent:命令执行的结果
    """
    shell.write(cmd)
    readcontent = shell.read()
    return readcontent

@mcp.tool()
def stop_shell()->bool:
    """
    在该函数调用之前请确保已经运行过get_a_shell这个函数
    请主动询问用户是否调用该函数
    如果调用该函数之后用户再次涉及到session或者后渗透相关内容则需要重新调用list_all_shell和get_a_shell函数

    Args:
        
    return:
        True:函数成功执行
    """
    shell.stop()
    return True


if __name__ == "__main__":
    # Initialize and run the server
    mcp.run(transport='stdio')